import nltk
from itertools import chain
from typing import List, Tuple

from nltk.translate.meteor_score import meteor_score

# nltk.download("wordnet")


def given_items_percentage(references: List[str], candidate: str) -> Tuple[int, int]:
    """Calculate the percentage of items in the candidate that are present in the references."""
    candidate_tokens = candidate.split()
    reference_tokens = list(chain.from_iterable([ref.split() for ref in references]))

    if not candidate_tokens:
        return 0

    count = 0
    for token in candidate_tokens:
        if token in reference_tokens:
            count += 1

    return count, len(candidate_tokens)


def corpus_given_items_percentage(
    references: List[List[str]], candidates: List[str]
) -> float:
    """Calculate the average percentage of items in the candidates that are present in the references."""
    total_count = 0
    total_length = 0
    for reference, candidate in zip(references, candidates):
        count, length = given_items_percentage(reference, candidate)
        total_count += count
        total_length += length

    return total_count / total_length


def extra_items_percentage(references: List[str], candidate: str) -> Tuple[int, int]:
    """Calculate the percentage of items in the candidate that are not present in the references."""
    candidate_tokens = candidate.split()
    reference_tokens = list(chain.from_iterable([ref.split() for ref in references]))

    if not candidate_tokens:
        return 0

    count = 0
    for token in candidate_tokens:
        if token not in reference_tokens:
            count += 1

    return count, len(candidate_tokens)


def corpus_extra_items_percentage(
    references: List[List[str]], candidates: List[str]
) -> float:
    """Calculate the average percentage of items in the candidates that are not present in the references."""
    total_count = 0
    total_length = 0
    for reference, candidate in zip(references, candidates):
        count, length = extra_items_percentage(reference, candidate)
        total_count += count
        total_length += length

    return total_count / total_length


def corpus_meteor(refs: List[List[str]], hyps: List[str]) -> float:
    """Calculate the METEOR score on a corpus
    Args:
        refs (list of list of str): List of reference translations where each reference translation is a list of tokens.
        hyps (list of str): List of hypothesis translations where each hypothesis translation is a list of tokens.

    Returns:
        float: The average METEOR score of the corpus.
    """
    scores = []
    for ref, hyp in zip(refs, hyps):
        # Note: `meteor_score` expects a list of reference sentences where each is a single string
        score = meteor_score([r.split() for r in ref], hyp.split())
        scores.append(score)
    return sum(scores) / len(scores) if scores else 0
